/*************************************************************************
 > File Name: monolith/interface/inference/utils.cc
 > Author: dzhang
 > Mail: dong.zhang@ingenic.com
 > Created Time : Sun 17 Jan 2021 10:44:34 PM CST
 > Description:
 ************************************************************************/

#include "monolith/interface/inference/utils.h"
#include <string.h>
#include <cmath>
CC_NAME_SPACE_BEGIN
namespace inference {

void nms(std::vector<Box> &input, std::vector<Box> &output,
        float iou_threshold = 0.5f, NMSType type = NMSType::hard_nms) {
    std::sort(input.begin(), input.end(), [](const Box &a, const Box &b) {
        return a.score > b.score;
    });

    int box_num = input.size();

    std::vector<int> merged(box_num, 0);

    for (int i = 0; i < box_num; i++) {
        if (merged[i])
            continue;
        std::vector<Box> buf;

        buf.push_back(input[i]);
        merged[i] = 1;

        float h0 = input[i].ry - input[i].ly + 1;
        float w0 = input[i].rx - input[i].lx + 1;

        float area0 = h0 * w0;

        for (int j = i + 1; j < box_num; j++) {
            if (merged[j])
                continue;

            float inner_x0 =
                    input[i].lx > input[j].lx ? input[i].lx : input[j].lx;
            float inner_y0 =
                    input[i].ly > input[j].ly ? input[i].ly : input[j].ly;

            float inner_lx =
                    input[i].rx < input[j].rx ? input[i].rx : input[j].rx;
            float inner_ly =
                    input[i].ry < input[j].ry ? input[i].ry : input[j].ry;

            float inner_h = inner_ly - inner_y0 + 1;
            float inner_w = inner_lx - inner_x0 + 1;

            if (inner_h <= 0 || inner_w <= 0)
                continue;

            float inner_area = inner_h * inner_w;

            float h1 = input[j].ry - input[j].ly + 1;
            float w1 = input[j].rx - input[j].lx + 1;

            float area1 = h1 * w1;

            float score;

            score = inner_area / (area0 + area1 - inner_area);

            if (score > iou_threshold) {
                merged[j] = 1;
                buf.push_back(input[j]);
            }
        }
        switch (type) {
        case NMSType::hard_nms: {
            output.push_back(buf[0]);
            break;
        }
        case NMSType::blending_nms: {
            float total = 0;
            for (int i = 0; i < buf.size(); i++) {
                total += exp(buf[i].score);
            }
            Box rects;
            memset(&rects, 0, sizeof(rects));
            for (int i = 0; i < buf.size(); i++) {
                float rate = exp(buf[i].score) / total;
                rects.lx += buf[i].lx * rate;
                rects.ly += buf[i].ly * rate;
                rects.rx += buf[i].rx * rate;
                rects.ry += buf[i].ry * rate;
                rects.score += buf[i].score * rate;
            }
            output.push_back(rects);
            break;
        }
        default: {
            printf("wrong type of nms.");
            exit(-1);
        }
        }
    }
}

int l2normal(std::vector<float> &data) {
    float nor = 0.0f;
    for (auto &tmp : data) {
        nor += pow(tmp, 2);
    }
    nor = sqrt(nor);
    for (int i = 0; i < data.size(); i++) {
        data[i] = data[i] / nor;
    }
    return 0;
}

float distance(std::vector<float> &a, std::vector<float> &b) {
    float dis = 0.0f;
    for (int i = 0; i < a.size(); i++) {
        dis += pow(a[i] - b[i], 2);
    }
    dis = sqrt(dis);
    return dis;
}

int crop_img(cv::Mat& src_img, cv::Mat& dst_img, int x, int y, int width,
        int height) {
    cv::Rect area(x, y, width, height);
    dst_img = src_img(area);
    return 0;
}

int crop_face_img(cv::Mat& src_img, cv::Mat& dst_img, int x, int y, int width,
        int height){
    int img_height = src_img.rows;
    int img_width = src_img.cols;
    float cx = x + width / 2.0f;
    float cy = y + height / 2.0f;
    float new_width = width + width * 0.;
    float new_height = height + height * 0.;
    int size = std::max(new_height, new_width);
    new_width = new_height = size;
    int lx = static_cast<int>(cx - new_width / 2.0);
    int ly = static_cast<int>(cy - new_height / 2.0);
    if(lx < 0 || ly < 0 || lx+new_width > img_width || ly+new_height > img_height){
        return 1;
    }
    lx = std::max(lx, 0);
    ly = std::max(ly, 0);
    return crop_img(src_img, dst_img, lx, ly,  static_cast<int>(new_width),  static_cast<int>(new_height));
}

} // namespace inference
CC_NAME_SPACE_END
